#!/bin/bash

# 脚本名称：convert_disk.sh

# 功能说明：
# 本脚本用于转换磁盘格式，并验证输入参数和路径有效性。

# 主要流程包括：
# 1. 安装必要的依赖工具（如 qemu-img 等）。
# 2. 解析命令行参数。
# 3. 验证输入参数。
# 4. 检查路径有效性。
# 5. 检查物理磁盘路径。
# 6. 执行磁盘转换。

# 注意事项：
# - 需要 root 权限运行部分命令，如 qemu-img。
# - 在执行磁盘转换前，建议备份重要数据，以防万一。

# 显示用法
usage() {
    echo "Usage: $0 [-i <input>] [-f <format>] [-o <output>] [--help]"
    echo "Options:"
    echo "  -i, --input     Path to the input disk file."
    echo "  -f, --format    Target format (e.g., qcow2, raw, vmdk)."
    echo "  -o, --output    Path to save the converted file or target physical disk."
    echo "  -h, --help      Display this help message."
    echo
    echo "Supported formats: qcow2, raw, vmdk, vdi, vpc, qed"
    exit 0
}

# 检查命令是否存在
ensure_command() {
    local cmd="$1"
    local package="$2"
    if ! command -v "$cmd" &>/dev/null; then
        echo "Installing $package..."
        sudo apt-get update && sudo apt-get install -y "$package"
        if ! command -v "$cmd" &>/dev/null; then
            echo "Error: Failed to install $package. Aborting."
            exit 1
        fi
    fi
}

# 解析命令行参数
parse_args() {
    while getopts "i:f:o:h" opt; do
        case "$opt" in
            i) input="$OPTARG" ;;
            f) format="$OPTARG" ;;
            o) output="$OPTARG" ;;
            h) usage ;;
            *) usage ;;
        esac
    done
}

# 检查依赖项
check_dependencies() {
    ensure_command "qemu-img" "qemu-utils"
}

# 检查路径有效性
validate_paths() {
    # 检查输入文件是否存在
    if [[ ! -e "$input" ]]; then
        echo "Error: Input file '$input' does not exist."
        exit 1
    fi

    # 如果输出路径是目录，自动生成文件名
    if [[ -d "$output" ]]; then
        output="${output}/$(basename "$input" .$format).$format"
    fi

    # 检查输出路径是否可写
    if [[ ! -w "$(dirname "$output")" && ! "$output" =~ ^/dev/ ]]; then
        echo "Error: Output directory is not writable: $(dirname "$output")"
        exit 1
    fi
}

# 检查物理磁盘路径
check_physical_disk() {
    if [[ "$output" == /dev/* ]]; then
        # 确保输出路径不是当前系统的根磁盘
        root_partition=$(findmnt / -o SOURCE -n)
        root_disk=$(lsblk -no PKNAME "$root_partition" 2>/dev/null | head -n 1)
        root_disk="/dev/${root_disk}"

        if [[ -z "$root_disk" ]]; then
            echo "Error: Unable to determine the root disk from the root partition '$root_partition'."
            exit 1
        fi

        if [[ "$output" == "$root_disk" || "$output" == "$root_partition" ]]; then
            echo "Error: Output disk '$output' is the root filesystem or its parent disk. Aborting to prevent system damage."
            exit 1
        fi

        # 提示用户确认写入物理磁盘
        echo "Warning: You are about to overwrite a physical disk ($output). Proceed with caution."
        read -p "Do you want to continue? (y/n): " confirmation
        if [[ "$confirmation" != "y" ]]; then
            echo "Operation aborted."
            exit 0
        fi
    fi
}

# 验证输入参数
validate_inputs() {
    if [[ -z "$input" || -z "$format" || -z "$output" ]]; then
        echo "Error: Missing required parameters."
        usage
    fi

    # 定义允许的格式列表
    local valid_formats_list=("qcow2" "raw" "vmdk" "vdi" "vpc" "qed")
    local format_valid=false

    # 验证格式是否在允许的列表中
    for valid_format in "${valid_formats_list[@]}"; do
        if [[ "$format" == "$valid_format" ]]; then
            format_valid=true
            break
        fi
    done

    if [[ "$format_valid" == false ]]; then
        echo "Error: Unsupported format '$format'. Supported formats are: ${valid_formats_list[*]}"
        exit 1
    fi
}

# 执行磁盘转换
convert_disk() {
    echo "Converting '$input' to '$format' format..."
    if ! sudo qemu-img convert -p -O "$format" "$input" "$output"; then
        echo "Error: Conversion failed."
        exit 1
    fi
    sync
    echo "Conversion completed successfully. Output file: $output"
}

# 主函数
main() {
    parse_args "$@"
    check_dependencies
    validate_inputs
    validate_paths
    check_physical_disk
    convert_disk
}

main "$@"

